# -*- coding: utf-8 -*-
import numpy as np
from sklearn import tree
import matplotlib.pyplot as plt
import graphviz
# 处理数据
filename = './Wine_random.csv'
labelname = './label_wine_randomed.csv'
data = np.loadtxt(open(filename, "rb"), delimiter=",", skiprows=0)
label = np.loadtxt(open(labelname, "rb"), dtype=int, delimiter=",", skiprows=0)
#
'''
80%数据用于训练，20%数据用于测试
data.shape[0]就是获取这个数据集的行数，即样本的数量。
data[:int(data.shape[0] * 0.8)]：这个切片操作从data中选取从开始到计算出的索引（即int(data.shape[0] * 0.8)）的所有样本。这部分数据被用作训练集。
切片操作 data[int(data.shape[0] * 0.8):] 从这个计算出的索引位置开始，一直切片到 data 的末尾。
'''
data_train = data[:int(data.shape[0] * 0.8)]
label_train = label[:int(label.shape[0] * 0.8)]
data_test = data[int(data.shape[0] * 0.8):]
label_test = label[int(label.shape[0] * 0.8):]
# 训练模型
''' 
基尼系数是一种衡量数据集纯度的指标，表示从数据集中随机抽取两个样本，其类别标签不一致的概率。
在决策树中，选择基尼系数最小的特征作为划分节点。而信息增益则更侧重于纯度的提升，即在划分后数据的不确定性程度的降低。
这两种方法都可以作为特征选择的标准，但默认情况下，tree.DecisionTreeClassifier使用的是基尼系数。
'''
clf = tree.DecisionTreeClassifier()
clf.fit(data_train, label_train)
label_predict = clf.predict(data_test)

# 与真实标签比较
x = range(data_test.shape[0])
fig = plt.figure()
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
ax1.set_title('Predict cluster')
ax2.set_title('True cluster')
plt.xlabel('samples')
plt.ylabel('label')
ax1.scatter(x, label_predict, c=label_predict, marker='o')
ax2.scatter(x, label_test, c=label_test, marker='s')
plt.show()

'''new add '''
dot_data = tree.export_graphviz(clf, out_file=None,
                         filled=True, rounded=True,
                         special_characters=True)
graph = graphviz.Source(dot_data)
graph.render('CLF_cart_RF.dot', "img/", view=True)